package com.zszq.ms.bigmodeltest.model.qwenurbo;

import com.alibaba.dashscope.aigc.generation.Generation;
import com.alibaba.dashscope.aigc.generation.GenerationParam;
import com.alibaba.dashscope.aigc.generation.GenerationResult;
import com.alibaba.dashscope.common.Message;
import com.alibaba.dashscope.common.Role;
import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import com.zszq.ms.bigmodeltest.entity.AiBigmodelConfig;
import com.zszq.ms.bigmodeltest.enums.BigModelEnum;
import com.zszq.ms.bigmodeltest.model.BaseModel;
import org.apache.commons.lang3.StringUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import java.util.Collections;

/**
 * 通义千问qwen-urbo访问工具类
 *
 * @author 孟帅
 * @since 2025/3/31
 */
@Component
public class QwenUrboService implements BaseModel {

    @Autowired
    private Generation generation;

    public GenerationResult call(String content, AiBigmodelConfig aiBigmodelConfig) throws Exception {
        Message userMessage = Message.builder().role(Role.USER.getValue()).content(content).build();
        GenerationParam param = GenerationParam.builder()
                // 指定用于对话的通义千问模型名
                .model(aiBigmodelConfig.getCode()).messages(Collections.singletonList(userMessage)).resultFormat(GenerationParam.ResultFormat.MESSAGE).topK(50).temperature(0.8f)
                // 生成过程中核采样方法概率阈值，例如，取值为0.8时，仅保留概率加起来大于等于0.8的最可能token的最小集合作为候选集。
                .topP(0.8)
                // 取值范围为（0,1.0)，取值越大，生成的随机性越高；取值越低，生成的确定性越高。
                .seed(1234)
                // 阿里云控制台DataScope获取的api-key
                .apiKey(aiBigmodelConfig.getKey())
                // 启用互联网搜索，模型会将搜索结果作为文本生成过程中的参考信息，但模型会基于其内部逻辑“自行判断”是否使用互联网搜索结果。
                .enableSearch(true).build();
        return generation.call(param);
    }

    @Override
    public JSONObject resTransferJson(String content, AiBigmodelConfig aiBigmodelConfig) throws Exception {
        GenerationResult generationResult = call(content, aiBigmodelConfig);
        JSONObject res = new JSONObject();
        // 解析返回结果
        String jsonResponse = generationResult.getOutput().getChoices().get(0).getMessage().getContent();
        if (StringUtils.isNotEmpty(jsonResponse)) {
            int start = jsonResponse.indexOf("{");
            int end = jsonResponse.indexOf("}", start);
            if (start != -1 && end != -1) {
                String result = jsonResponse.substring(start, end + 1);
                res = JSON.parseObject(result);
            }
        }
        return res;
    }

    @Override
    public String getModelCode() {
        return BigModelEnum.QWEN_TURBO.getCode();
    }
}
